<?php
/**
 * Created by PhpStorm.
 * User: Administrator
 * Date: 2024/3/20
 * Time: 11:46
 */

namespace app\queue\aiimageadmin;

use app\common\lib\Common;
use app\common\lib\Constant;
use app\common\lib\Guzzle;
use app\common\lib\RedisQueue;
use app\exception\IndexException;
use app\model\queue\AiImageRequestQueueModel;
use app\model\redis\RedisQueueModel;
use app\model\user\UserNovelChapterPartModel;
use support\Redis;
use Webman\RedisQueue\Consumer;
use app\common\lib\TencentCos;

class RequestAiImageAdmin  implements Consumer
{
    // 要消费的队列名
    public $queue = 'requestAiImageAdmin';

    // 连接名，对应 plugin/webman/redis-queue/redis.php 里的连接`
    public $connection = 'ai_image_consumer';

    protected $resultData = [];


    protected function checkRequestId($data,$key)
    {
        //不在当前环境中的，冲压入队列尾部，然后第一个消费
        if(!isset($data['env']) || $data['env']!=config('app.app_env')){
            //抛出异常，进入错误队列中，发起重试
            //throw new IndexException(4020,'环境不对');
            $queueKey = '{redis-queue}-waiting'.$this->queue;
            $now = time();
            $package_str = json_encode([
                'id'       => $now.rand(),
                'time'     => $now,
                'delay'    => 0,
                'attempts' => 0,
                'queue'    => $this->queue,
                'data'     => $data
            ]);
            \Webman\RedisQueue\Redis::connection('ai_image_consumer')->rpush($queueKey,$package_str);

            while (true){
                //获取下一个元素
                $redisQueues = \Webman\RedisQueue\Redis::connection('ai_image_consumer')->lRange($queueKey,0,-1);
                if($redisQueues){
                    $redisQueues = array_reverse($redisQueues);
                    $qData = json_decode($redisQueues[0], true);
                    if(!empty($qData['data']['env']) && $qData['data']['env']!=config('app.app_env')){
                        sleep(3);
                        continue;
                    }
                }
                break;
            }
            return false;
        }

        //如果客户端生图队列还有排队，则优先生成客户端生图队列
        $userQueueKey = '{redis-queue}-waitingrequestAiImage';
        while (true){
            //获取下一个元素
            $redisQueues = \Webman\RedisQueue\Redis::connection('ai_image_consumer')->lRange($userQueueKey,0,-1);
            if($redisQueues==false){
                break;
            }
            sleep(1);
        }

        $aiImageServers = Redis::connection('ai_image_cache')->get('ai_image_server_manage');
        if(empty($aiImageServers)){
            $aiImageServers = Constant::AI_IMAGE_SERVER_DOMAIN_LIST;
        }else{
            $aiImageServers = json_decode($aiImageServers, true);
        }

        //检查上次执行队列是否已完成，超时略过
        $ipServer = [];
        while (true){
            foreach ($aiImageServers as $val){
                $cacheData = Redis::connection('ai_image_cache')->get($key.'_'.$val['ip']);
                $cacheData && $cacheData = json_decode($cacheData, true);
                if($cacheData==false){
                    $ipServer = $val;
                    break;
                }elseif(!empty($cacheData['start_time']) && (time()-$cacheData['start_time'])>180){
                    $ipServer = $val;
                    break;
                }
            }

            if(!empty($ipServer)){
                break;
            }

            sleep(1);
        }

        if($ipServer){
            $endServer = [];
            foreach ($aiImageServers as $k=>$val){
                if($val['ip']==$ipServer['ip']){
                    $endServer = $val;
                    unset($aiImageServers[$k]);
                }
            }
            array_push($aiImageServers,$endServer);
            Redis::connection('ai_image_cache')->set('ai_image_server_manage',json_encode(array_values($aiImageServers)));
        }

        return $ipServer;
    }

    // 消费
    public function consume($data)
    {
        $times = [10000,20000,30000,40000,50000];
        shuffle($times);
        usleep($times[0]);

        Common::initQueueConfig();

        //执行中加入缓存记录
        $key = 'ai_image_server_cache';
        $ipServer = $this->checkRequestId($data,$key);
        if($ipServer==false){
            return false;
        }

        $cacheData = array_merge($data,['ip'=>$ipServer['ip'],'start_time'=>time()]);
        Redis::connection('ai_image_cache')->set($key.'_'.$ipServer['ip'],json_encode($cacheData,256));

        try{
            $queue = RedisQueueModel::getInstance()->find($data['queueid']);
            if($queue==false){
                throw new IndexException(4020,'队列不存在');
            }
            $queue->exec_time = date('Y-m-d H:i:s');
            $queue->save();

            // 无需反序列化
            $result = $this->handle($data,$ipServer);

            RedisQueue::getInstance()->saveQueue($this->queue,$data,$result);
        }catch (IndexException $e){
            $first = AiImageRequestQueueModel::getInstance()->getWhere(['id'=>$data['id']])->first();
            $first->status = 3;
            $first->result = json_encode(['error'=>$e->getMessage()],256);
            $first->end_time = date('Y-m-d H:i:s');
            $first->tts = round(time()-strtotime($first->start_time));
            $first->save();

            RedisQueue::getInstance()->saveQueue($this->queue,$data,['error'=>$e->getMessage()]);
            Common::queueExecptionReport($queue,$data,$e);
        }catch (\Throwable $e){
            $first = AiImageRequestQueueModel::getInstance()->getWhere(['id'=>$data['id']])->first();
            $first->status = 3;
            $first->result = json_encode(['error'=>$e->getMessage()],256);
            $first->end_time = date('Y-m-d H:i:s');
            $first->tts = round(time()-strtotime($first->start_time));
            $first->save();

            RedisQueue::getInstance()->saveQueue($this->queue,$data,['error'=>$e->getMessage()]);
            Common::queueExecptionReport($queue,$data,$e);
        }
        Redis::connection('ai_image_cache')->del($key.'_'.$ipServer['ip']);
    }

    protected function handle($data,$ipServer,$nums=1)
    {
        //todo  生成ai图像
        $first = AiImageRequestQueueModel::getInstance()->getWhere(['id'=>$data['id']])->first();
        if($first==false){
            throw new IndexException(4020,$data['id'].' 请求生成ai图像不存在');
        }

        $first->status = 1;
        $first->ip = $ipServer['ip']??'';
        $first->start_time = date('Y-m-d H:i:s');
        $first->save();

        if(strpos($first['source'],'user_plan_part_image_save_')!==false){
            $sources = explode('_',$first['source']);
            $partId = $sources[count($sources)-1];
            $userPart = UserNovelChapterPartModel::getInstance()->getWhere(['id'=>$partId])->first();
            if($userPart && !empty($userPart['ai_images_arr'])){
                $result = ['images'=>json_decode($userPart['ai_images_arr'], true)];
                $first->status = 2;
                $first->result = json_encode($result,256);
                $first->tts = round(time()-strtotime($first->start_time));
                $first->end_time = date('Y-m-d H:i:s');
                $first->save();

                return ['生成成功'];
            }
        }

        $param = json_decode($first['param'],true);
        if($param==false){
            $first->status = 3;
            $first->result = json_encode(['error'=>'指令参数不能为空或者错误'],256);
            $first->end_time = date('Y-m-d H:i:s');
            $first->tts = round(time()-strtotime($first->start_time));
            $first->save();

            return ['指令参数不能为空或者错误'];
        }
        $headers = [
            'Authorization'=>'Basic dHVpd2VuOlJrSWopWCdZYWh6TGRIcUI=',
            'content-type'=>'application/json'
        ];

        if($param['batch_size']>3){
            Guzzle::setTimeOut(240);
        }else{
            Guzzle::setTimeOut(180);
        }
        if(isset($param['alwayson_scripts']['controlnet']['args'][0]) && is_file(runtime_path('input_image.png'))){
            $param['alwayson_scripts']['controlnet']['args'][0]['image'] = base64_encode(file_get_contents(runtime_path('input_image.png')));
        }
        if(isset($param['alwayson_scripts']['controlnet']['args'][0]['model'])){
            unset($param['alwayson_scripts']['controlnet']['args'][0]['model']);
        }
        //todo 根据这个参数判断是否为图生图
        if(!empty($param['init_images'])){
            $initImages = [];
            foreach ($param['init_images'] as $img){
                $initImages[] = base64_encode(file_get_contents($img));
            }
            $param['init_images'] = $initImages;

            $apiUrl = 'sdapi/v1/img2img';
            $result = Guzzle::post($ipServer['domain'],$apiUrl,$param,$headers,'json');
        }else{
            $apiUrl = 'sdapi/v1/txt2img';
            $result = Guzzle::post($ipServer['domain'],$apiUrl,$param,$headers,'json');
        }
        //file_put_contents(runtime_path('request_image_param'),json_encode($param,256));
        if(!empty($result['images'])){
            if(is_array($result['images'])){
                $images = [];
                foreach ($result['images'] as $k=>$image){
                    if($first['source']=='admin_batch_image_test'){
                        $filePath = 'aiimage/'.date('ymd').'/'.$data['requestid'].'_'.$k.'.png';
                        Common::mkdirs(dirname(public_path($filePath)));
                        $fileName = public_path($filePath);
                        file_put_contents($fileName,file_get_contents('data:image/png;base64,'.$image));
                        $k==0 && $images[] = rtrim(config('app.app_url'),'/').'/'.$filePath;
                        $k==1 && $images[] = rtrim(config('app.app_url'),'/').'/'.$filePath;
                        $k==2 && $images[] = rtrim(config('app.app_url'),'/').'/'.$filePath;
                        $k==3 && $images[] = rtrim(config('app.app_url'),'/').'/'.$filePath;
                    }else{
                        $fileName = runtime_path('stable_diffusion_'.$data['requestid'].'_'.$k.'.png');
                        file_put_contents($fileName,file_get_contents('data:image/png;base64,'.$image));
                        $res = TencentCos::getInstance()->uploadLocalFile($fileName,'stable_diffusion_image');

                        $k==0 && $images[] ='https://'.$res['Location'];
                        $k==1 && $images[] ='https://'.$res['Location'];
                        $k==2 && $images[] ='https://'.$res['Location'];
                        $k==3 && $images[] ='https://'.$res['Location'];

                        @unlink($fileName);
                    }
                }
                $result['images'] = $images;

                if(isset($result['parameters']['alwayson_scripts']['controlnet']['args'][0]['input_image'])){
                    unset($result['parameters']['alwayson_scripts']['controlnet']['args'][0]['input_image']);
                }
                if(isset($result['info'])){
                    //unset($result['info']);
                }
                $result['api_url'] = $apiUrl;

                $first->status = 2;
                $first->result = json_encode($result,256);
                $first->end_time = date('Y-m-d H:i:s');
                $first->tts = round(time()-strtotime($first->start_time));
                $first->save();

                return ['生成成功'];
            }else{
                if(isset($result['parameters']['alwayson_scripts']['controlnet']['args'][0]['input_image'])){
                    unset($result['parameters']['alwayson_scripts']['controlnet']['args'][0]['input_image']);
                }
                if(isset($result['info'])){
                    //unset($result['info']);
                }
                $result['api_url'] = $apiUrl;

                $first->status = 3;
                $first->result = json_encode($result,256);
                $first->end_time = date('Y-m-d H:i:s');
                $first->tts = round(time()-strtotime($first->start_time));
                $first->save();

                return ['ai生图失败'];
            }
        }

        if(strpos(strtolower(Guzzle::getError()),'curl')!==false){
            if($nums==1){
                //重试一次
               return $this->handle($data,$ipServer,2);
            }
            $first->status = 4;
        }else{
            $first->status = 3;
        }

        $first->result = json_encode(['error'=>Guzzle::getError()??'请求错误'],256);
        $first->end_time = date('Y-m-d H:i:s');
        $first->tts = round(time()-strtotime($first->start_time));
        $first->save();

        throw new IndexException(4020,(Guzzle::getError()??'请求失败,ai生图失败'));
    }

    // 消费失败回调
    /*
    $package = [
        'id' => 1357277951, // 消息ID
        'time' => 1709170510, // 消息时间
        'delay' => 0, // 延迟时间
        'attempts' => 2, // 消费次数
        'queue' => 'send-mail', // 队列名
        'data' => ['to' => 'tom@gmail.com', 'content' => 'hello'], // 消息内容
        'max_attempts' => 5, // 最大重试次数
        'error' => '错误信息' // 错误信息
    ]
    */
    public function onConsumeFailure(\Throwable $e, $package)
    {
//        echo "consume failure\n";
//        echo $e->getMessage() . "\n";
//        // 无需反序列化
//        var_export($package);

        //RedisQueue::getInstance()->saveQueue($this->queue,$package['data'],['error'=>$e->getMessage()]);

        $package['attempts']=1;
        return $package;
    }

}